import numpy as np
import pywt
import cv2
import matplotlib.pyplot as plt

# 归一化数据
def normalization(data):
    range = np.max(data) - np.min(data)
    return (data - np.min(data)) / range

img = cv2.imread("beauty.jpg", cv2.IMREAD_GRAYSCALE).astype(np.float32)
# 二维小波分解
coeffs = pywt.wavedec2(img, 'haar', level=3)

# 合并系数
cA = normalization(coeffs[0])
for i in range(1, len(coeffs)):
    cH = normalization(coeffs[i][0])
    cV = normalization(coeffs[i][1])
    cD = normalization(coeffs[i][2])
    AH = np.concatenate([cA, cH], axis=1)
    VD = np.concatenate([cV, cD], axis=1)
    cA = np.concatenate([AH, VD], axis=0)
# 显示分解后小波系数
plt.imshow(cA, 'gray')
plt.title('2D WT')
plt.show()
